{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e8952066-7a10-4c9b-a4b7-27be074ae269",
   "metadata": {},
   "source": [
    "## Create Keycloak resources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7252812d-90eb-4752-91a7-d46b400bacd8",
   "metadata": {},
   "source": [
    "Wait until Keycloak is running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e5d13f76-f184-44f6-8542-54a61060e531",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\u001b[36m\"Status\"\u001b[0m:\u001b[32m \"running\"\u001b[0m,\u001b[36m \"Running\"\u001b[0m:\u001b[95m true\u001b[0m,\u001b[36m \"Paused\"\u001b[0m:\u001b[95m false\u001b[0m,\u001b[36m \"Restarting\"\u001b[0m:\u001b[95m false\u001b[0m,\u001b[36m \"OOMKilled\"\u001b[0m:\u001b[95m false\u001b[0m,\u001b[36m \"Dead\"\u001b[0m:\u001b[95m false\u001b[0m,\u001b[36m \"Pid\"\u001b[0m:\u001b[95m 2838024\u001b[0m,\u001b[36m \"ExitCode\"\u001b[0m:\u001b[95m 0\u001b[0m,\u001b[36m \"Error\"\u001b[0m:\u001b[32m \"\"\u001b[0m,\u001b[36m \"StartedAt\"\u001b[0m:\u001b[32m \"2024-09-09T06:37:49.055739669Z\"\u001b[0m,\u001b[36m \"FinishedAt\"\u001b[0m:\u001b[32m \"0001-01-01T00:00:00Z\"\u001b[0m}\n"
     ]
    }
   ],
   "source": [
    "!docker inspect --format='json' my-keycloak | yq '.[0].State'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc9a6329-4e89-464c-ac48-dbbadaf72a2b",
   "metadata": {},
   "source": [
    "Then create a sample realm and client with some roles and users matching the test environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d5c60591-f41d-4a5e-9b18-93385a889495",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "from dotenv import set_key\n",
    "\n",
    "OIDC_SERVER_URL = \"http://0.0.0.0:9999\"\n",
    "ADMIN_USERNAME = \"admin\"\n",
    "ADMIN_PASSWORD = \"admin\"\n",
    "\n",
    "access_token: str = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d16969bc-423a-4d18-afa3-97a791b84b13",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_token():\n",
    "    token_url = f\"{OIDC_SERVER_URL}/realms/master/protocol/openid-connect/token\"\n",
    "\n",
    "    token_data = {\n",
    "        \"grant_type\": \"password\",\n",
    "        \"client_id\": \"admin-cli\",\n",
    "        \"username\": ADMIN_USERNAME,\n",
    "        \"password\": ADMIN_PASSWORD,\n",
    "    }\n",
    "\n",
    "    token_response = requests.post(token_url, data=token_data)\n",
    "    if token_response.status_code == 200:\n",
    "        global access_token\n",
    "        access_token = token_response.json()[\"access_token\"]\n",
    "        return access_token\n",
    "    else:\n",
    "        print(\n",
    "            f\"Failed to obtain access token: {token_response.status_code} - {token_response.text}\"\n",
    "        )\n",
    "        raise Exception(\"Not authenticated\")\n",
    "\n",
    "\n",
    "def keycloak_post(endpoint, data=None):\n",
    "    url = f\"{OIDC_SERVER_URL}/admin/{endpoint}\"\n",
    "    print(f\"Creating {endpoint}\")\n",
    "    global access_token\n",
    "    headers = {\n",
    "        \"Content-Type\": \"application/json\",\n",
    "        \"Authorization\": f\"Bearer {access_token}\",\n",
    "    }\n",
    "    response = requests.request(\"POST\", url, headers=headers, data=json.dumps(data))\n",
    "    print(f\"POST response.status_code is {response.status_code}\")\n",
    "    return response.status_code\n",
    "\n",
    "\n",
    "def keycloak_get(endpoint):\n",
    "    url = f\"{OIDC_SERVER_URL}/admin/{endpoint}\"\n",
    "    global access_token\n",
    "    headers = {\n",
    "        \"Content-Type\": \"application/json\",\n",
    "        \"Authorization\": f\"Bearer {access_token}\",\n",
    "    }\n",
    "    response = requests.request(\"GET\", url, headers=headers)\n",
    "    print(f\"GET response.status_code is {response.status_code}\")\n",
    "    return response.json()\n",
    "\n",
    "\n",
    "def create_realm(realm_name):\n",
    "    data = {\"realm\": realm_name, \"enabled\": \"true\"}\n",
    "    keycloak_post(\"realms\", data=data)\n",
    "    response = keycloak_get(f\"realms/{realm_name}\")\n",
    "    return response[\"id\"]\n",
    "\n",
    "\n",
    "def create_client(realm_name, client_name):\n",
    "    data = {\n",
    "        \"clientId\": client_name,\n",
    "        \"enabled\": \"true\",\n",
    "        \"redirectUris\": [\n",
    "            \"http://localhost:8000/*\",\n",
    "            \"http://127.0.0.1:8000/*\",\n",
    "            \"http://0.0.0.0:8000/*\",\n",
    "        ],\n",
    "        \"publicClient\": False,\n",
    "        \"authorizationServicesEnabled\": True,\n",
    "        \"protocol\": \"openid-connect\",\n",
    "        \"standardFlowEnabled\": True,\n",
    "        \"directAccessGrantsEnabled\": True,\n",
    "        \"serviceAccountsEnabled\": True,\n",
    "    }\n",
    "    keycloak_post(f\"realms/{realm_name}/clients\", data=data)\n",
    "    response = keycloak_get(f\"realms/{realm_name}/clients\")\n",
    "    client = None\n",
    "    for c in response:\n",
    "        if c[\"clientId\"] == client_name:\n",
    "            client = c\n",
    "            break\n",
    "    client_id = client[\"id\"]\n",
    "    client_secret = client[\"secret\"]\n",
    "    return client_id, client_secret\n",
    "\n",
    "\n",
    "def create_client_roles(realm_name, client_id, roles):\n",
    "    for role_name in roles:\n",
    "        data = {\"name\": role_name, \"clientRole\": True}\n",
    "        keycloak_post(f\"realms/{realm_name}/clients/{client_id}/roles\", data=data)\n",
    "\n",
    "    response = keycloak_get(f\"realms/{realm_name}/clients/{client_id}/roles\")\n",
    "    roles_by_name = dict((role[\"name\"], role[\"id\"]) for role in response)\n",
    "    print(roles_by_name)\n",
    "    return roles_by_name\n",
    "\n",
    "\n",
    "def create_user_with_roles(\n",
    "    realm_name, username, password, client_id, roles_by_name, roles\n",
    "):\n",
    "    data = {\n",
    "        \"username\": username,\n",
    "        \"enabled\": True,\n",
    "        \"email\": f\"{username}@poc.com\",\n",
    "        \"emailVerified\": True,\n",
    "        \"firstName\": \"user\",\n",
    "        \"lastName\": f\"{username}\",\n",
    "        \"credentials\": [{\"type\": \"password\", \"value\": password}],\n",
    "        \"realmRoles\": [],\n",
    "    }\n",
    "    keycloak_post(f\"realms/{realm_name}/users\", data=data)\n",
    "    response = keycloak_get(f\"realms/{realm_name}/users\")\n",
    "    user = None\n",
    "    for u in response:\n",
    "        if u[\"username\"] == username:\n",
    "            user = u\n",
    "            break\n",
    "    user_id = user[\"id\"]\n",
    "\n",
    "    data = [\n",
    "        {\n",
    "            \"id\": roles_by_name[role_name],\n",
    "            \"name\": role_name,\n",
    "        }\n",
    "        for role_name in roles\n",
    "    ]\n",
    "    keycloak_post(\n",
    "        f\"realms/{realm_name}/users/{user_id}/role-mappings/clients/{client_id}\",\n",
    "        data=data,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e097fac1-f2c3-4afe-b78c-2c8279e3a84e",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_token()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f114fa41-8cea-486f-baf4-998cbf69fea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "realm_name = \"rbac_example\"\n",
    "client_name = \"app\"\n",
    "password = \"password\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0f889548-9b60-448b-beed-ac3fc1890b13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating realms\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "Creating realms/rbac_example/clients\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "Creating realms/rbac_example/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b/roles\n",
      "POST response.status_code is 201\n",
      "Creating realms/rbac_example/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b/roles\n",
      "POST response.status_code is 201\n",
      "Creating realms/rbac_example/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b/roles\n",
      "POST response.status_code is 201\n",
      "Creating realms/rbac_example/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b/roles\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "{'store_admin': '2d7a675f-031d-42b1-aba6-eb28a95561af', 'batch_admin': '8664084a-4e3c-42b0-8e37-70a8fea012b3', 'reader': '6cbf4473-c165-48bd-b572-d20133ae2b2b', 'uma_protection': '172d464d-92c7-4055-95af-3e048d8077b2', 'fresh_writer': '9e2abf47-a7af-414e-bf14-2c9897933532'}\n"
     ]
    }
   ],
   "source": [
    "realm_id = create_realm(realm_name)\n",
    "client_id, client_secret = create_client(realm_name, client_name)\n",
    "\n",
    "roles_by_name = create_client_roles(\n",
    "    realm_name, client_id, [\"reader\", \"fresh_writer\", \"store_admin\", \"batch_admin\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a3430d83-107d-44ad-acf2-0df810dff0ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating realms/rbac_example/users\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "Creating realms/rbac_example/users/a87b4ca8-e1a9-40f7-a166-f48fe45beec2/role-mappings/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b\n",
      "POST response.status_code is 204\n",
      "Creating realms/rbac_example/users\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "Creating realms/rbac_example/users/eb343a9b-d800-4fff-96b6-4588c7db08de/role-mappings/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b\n",
      "POST response.status_code is 204\n",
      "Creating realms/rbac_example/users\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "Creating realms/rbac_example/users/91bfbaae-e1fd-4167-9432-2d1d8ca8c838/role-mappings/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b\n",
      "POST response.status_code is 204\n",
      "Creating realms/rbac_example/users\n",
      "POST response.status_code is 201\n",
      "GET response.status_code is 200\n",
      "Creating realms/rbac_example/users/4d67e8ca-6c2a-48b7-b511-c3f6197aa5ae/role-mappings/clients/c3475e89-27c3-41ac-a3d1-0bbcaf68083b\n",
      "POST response.status_code is 204\n"
     ]
    }
   ],
   "source": [
    "create_user_with_roles(\n",
    "    realm_name, \"reader\", password, client_id, roles_by_name, [\"reader\"]\n",
    ")\n",
    "create_user_with_roles(\n",
    "    realm_name,\n",
    "    \"writer\",\n",
    "    password,\n",
    "    client_id,\n",
    "    roles_by_name,\n",
    "    [\"fresh_writer\"],\n",
    ")\n",
    "create_user_with_roles(\n",
    "    realm_name,\n",
    "    \"batch_admin\",\n",
    "    password,\n",
    "    client_id,\n",
    "    roles_by_name,\n",
    "    [\"batch_admin\"],\n",
    ")\n",
    "create_user_with_roles(\n",
    "    realm_name,\n",
    "    \"admin\",\n",
    "    password,\n",
    "    client_id,\n",
    "    roles_by_name,\n",
    "    [\"store_admin\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "54317f9e-476b-4b8e-864a-a07c54b549f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Realm rbac_example setup completed.\n",
      "Client app created with ID c3475e89-27c3-41ac-a3d1-0bbcaf68083b and secret REDACTED\n",
      "Settings configured in .env\n"
     ]
    }
   ],
   "source": [
    "print(f\"Realm {realm_name} setup completed.\")\n",
    "print(\n",
    "    f\"Client {client_name} created with ID {client_id} and secret {client_secret}\"\n",
    ")\n",
    "\n",
    "env_file = \".env\"\n",
    "with open(env_file, \"w\") as file:\n",
    "    pass\n",
    "\n",
    "# Write property P=1 to the .env file\n",
    "set_key(env_file, \"OIDC_SERVER_URL\", OIDC_SERVER_URL)\n",
    "set_key(env_file, \"REALM\", realm_name)\n",
    "set_key(env_file, \"CLIENT_ID\", client_name)\n",
    "set_key(env_file, \"CLIENT_SECRET\", client_secret)\n",
    "set_key(env_file, \"PASSWORD\", password)\n",
    "print(f\"Settings configured in {env_file}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35dcd5ed-4004-4570-965f-0f68668605d8",
   "metadata": {},
   "source": [
    "The [.env](.env) file contains the settings of the created realm, including the client secret to be used to connect the server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "46a1e2c7-e379-461d-b0bf-82354378e830",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OIDC_SERVER_URL='http://0.0.0.0:9999'\n",
      "REALM='rbac_example'\n",
      "CLIENT_ID='app'\n",
      "CLIENT_SECRET='REDACTED'\n",
      "PASSWORD='password'\n"
     ]
    }
   ],
   "source": [
    "!cat .env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d035826b-64d6-47cc-a48e-26eb29b31fc7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
